use std::{
    f64::consts::{LN_10, LN_2}, ops::{Add, Div, Mul, Neg, Sub}, ptr::NonNull
};

enum NodeType {
    Variable,
    AddX(usize, usize),
    SubX(usize, usize),
    MulX(usize, usize),
    DivX(usize, usize),
    Add(usize, f64),
    Sub(usize, f64),
    Mul(usize, f64),
    Div(usize, f64),
    // Rem(usize, usize),
    Neg(usize),
    Recip(usize),

    Sin(usize),
    Asin(usize),
    Sinh(usize),
    Asinh(usize),
    Cos(usize),
    Acos(usize),
    Cosh(usize),
    Acosh(usize),
    Tan(usize),
    Atan(usize),
    Tanh(usize),
    Atanh(usize),
    // ATan2(usize),
    Exp(usize),
    Exp2(usize),  // 2^x
    ExpM1(usize), // e^x - 1
    Ln(usize),
    Log10(usize),
    Log2(usize),
    Log(usize, f64),
    LogX(usize, usize),
    Powf(usize, f64),
    Powi(usize, i32),
    PowX(usize, usize),

    Sqrt(usize),
    Cbrt(usize),
    Hypot(usize, usize),
}
struct Node {
    node_type: NodeType,
    value: f64,
    grad: f64,
}

pub struct Graph<const N: usize> {
    nodes: Vec<Node>,
}

#[derive(Clone, Copy)]
pub struct Variable<const N: usize> {
    graph: NonNull<Graph<N>>,
    id: usize,
}

impl<const N: usize> Graph<N> {
    pub fn new() -> Self {
        Self { nodes: Vec::new() }
    }

    pub fn new_variable(&mut self, value: f64) -> Variable<N> {
        let id = self.nodes.len();
        let node = Node {
            node_type: NodeType::Variable,
            value,
            grad: 0.0,
        };
        self.nodes.push(node);
        Variable {
            graph: NonNull::from(self),
            id,
        }
    }

    pub fn grad(&self, x: Variable<N>) -> f64 {
        self.nodes[x.id].grad
    }
    pub fn value(&self, f: Variable<N>) -> f64 {
        self.nodes[f.id].value
    }

    pub fn backward(&mut self, output: Variable<N>) {
        self.nodes[output.id].grad = 1.0;
        for i in (0..self.nodes.len()).rev() {
            let grad_out = self.nodes[i].grad;
            if grad_out == 0.0 {
                continue;
            }

            match self.nodes[i].node_type {
                NodeType::Variable => {}
                NodeType::AddX(l_id, r_id) => {
                    self.nodes[l_id].grad += grad_out;
                    self.nodes[r_id].grad += grad_out;
                }
                NodeType::Add(l_id, _rhs) => {
                    self.nodes[l_id].grad += grad_out;
                }
                NodeType::SubX(l_id, r_id) => {
                    self.nodes[l_id].grad -= grad_out;
                    self.nodes[r_id].grad -= grad_out;
                }
                NodeType::Sub(l_id, _rhs) => {
                    self.nodes[l_id].grad -= grad_out;
                }
                NodeType::MulX(l_id, r_id) => {
                    self.nodes[l_id].grad += self.nodes[r_id].value * grad_out;
                    self.nodes[r_id].grad += self.nodes[l_id].value * grad_out;
                }
                NodeType::Mul(l_id, rhs) => {
                    self.nodes[l_id].grad += rhs * grad_out;
                }
                NodeType::DivX(l_id, r_id) => {
                    let denom = self.nodes[r_id].value;
                    self.nodes[l_id].grad += grad_out / denom;
                    self.nodes[r_id].grad -= grad_out * self.nodes[l_id].value / (denom * denom);
                }
                NodeType::Div(l_id, rhs) => {
                    self.nodes[l_id].grad += grad_out / rhs;
                }
                NodeType::Neg(id) => {
                    self.nodes[id].grad -= grad_out;
                }
                NodeType::Recip(id) => {
                    let denom = self.nodes[id].value;
                    self.nodes[id].grad -= grad_out / (denom * denom);
                }

                NodeType::Sin(id) => {
                    self.nodes[id].grad += grad_out * self.nodes[id].value.cos();
                }
                NodeType::Asin(id) => {
                    self.nodes[id].grad += grad_out / (1.0 - self.nodes[id].value.powi(2)).sqrt();
                }
                NodeType::Sinh(id) => {
                    self.nodes[id].grad += grad_out * self.nodes[id].value.cosh();
                }
                NodeType::Asinh(id) => {
                    self.nodes[id].grad += grad_out / self.nodes[id].value.powi(2).add(1.0).sqrt();
                }
                NodeType::Cos(id) => {
                    self.nodes[id].grad -= grad_out * self.nodes[id].value.sin();
                }
                NodeType::Acos(id) => {
                    self.nodes[id].grad -= grad_out / (1.0 - self.nodes[id].value.powi(2)).sqrt();
                }
                NodeType::Cosh(id) => {
                    self.nodes[id].grad += grad_out * self.nodes[id].value.sinh();
                }
                NodeType::Acosh(id) => {
                    self.nodes[id].grad += grad_out / self.nodes[id].value.powi(2).sub(1.0).sqrt();
                }
                NodeType::Tan(id) => {
                    self.nodes[id].grad += grad_out / self.nodes[id].value.cos().powi(2);
                }
                NodeType::Atan(id) => {
                    self.nodes[id].grad += grad_out / (1.0 + self.nodes[id].value.powi(2));
                }
                NodeType::Tanh(id) => {
                    self.nodes[id].grad += grad_out / self.nodes[id].value.cosh().powi(2);
                }
                NodeType::Atanh(id) => {
                    self.nodes[id].grad += grad_out / (1.0 - self.nodes[id].value.powi(2));
                }

                NodeType::Exp(id) => {
                    self.nodes[id].grad += grad_out * self.nodes[id].value.exp();
                }
                NodeType::Exp2(id) => {
                    self.nodes[id].grad +=
                        grad_out * self.nodes[id].value.exp2() * LN_2;
                }
                NodeType::ExpM1(id) => {
                    self.nodes[id].grad +=
                        grad_out * self.nodes[id].value.exp2();
                }
                NodeType::Ln(id) => {
                    self.nodes[id].grad += grad_out / self.nodes[id].value;
                }
                NodeType::Log10(id) => {
                    self.nodes[id].grad += grad_out / (self.nodes[id].value * LN_10);
                }
                NodeType::Log2(id) => {
                    self.nodes[id].grad += grad_out / (self.nodes[id].value * LN_2);
                }
                NodeType::Log(id, base) => {
                    self.nodes[id].grad += grad_out / (self.nodes[id].value * base.ln());
                }
                NodeType::LogX(l_id, r_id) => {
                    self.nodes[l_id].grad +=
                        grad_out / (self.nodes[l_id].value * self.nodes[r_id].value.ln());
                    self.nodes[r_id].grad -= grad_out
                        * self.nodes[l_id].value.log(self.nodes[r_id].value)
                        / (self.nodes[r_id].value * self.nodes[r_id].value.ln());
                }
                NodeType::Powf(id, n) => {
                    self.nodes[id].grad += grad_out * n * self.nodes[id].value.powf(n - 1.0);
                }
                NodeType::Powi(id, n) => {
                    self.nodes[id].grad += grad_out * n as f64 * self.nodes[id].value.powi(n - 1);
                }
                NodeType::PowX(l_id, r_id) => {
                    self.nodes[l_id].grad += grad_out
                        * self.nodes[r_id].value
                        * self.nodes[l_id].value.powf(self.nodes[r_id].value - 1.0);
                    self.nodes[r_id].grad += grad_out
                        * self.nodes[l_id].value.powf(self.nodes[r_id].value)
                        * self.nodes[l_id].value.ln();
                }

                NodeType::Sqrt(id) => {
                    self.nodes[id].grad += grad_out / (2.0 * self.nodes[id].value.sqrt());
                }
                NodeType::Cbrt(id) => {
                    self.nodes[id].grad += grad_out / (3.0 * self.nodes[id].value.powf(2.0 / 3.0));
                }
                NodeType::Hypot(l_id, r_id) => {
                    let denom =
                        (self.nodes[l_id].value.powi(2) + self.nodes[r_id].value.powi(2)).sqrt();
                    self.nodes[l_id].grad += self.nodes[l_id].value / denom;
                    self.nodes[r_id].grad += self.nodes[r_id].value / denom;
                }
            }
        }
    }
}

impl<const N: usize> Add for Variable<N> {
    type Output = Self;
    fn add(self, rhs: Self) -> Self::Output {
        let mut lhs = self;
        let graph = unsafe { lhs.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::AddX(self.id, rhs.id),
            value: graph.nodes[self.id].value + graph.nodes[rhs.id].value,
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }
}
impl<const N: usize> Add<f64> for Variable<N> {
    type Output = Self;
    fn add(self, rhs: f64) -> Self::Output {
        let mut lhs = self;
        let graph = unsafe { lhs.graph.as_mut() };

        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Add(self.id, rhs),
            value: graph.nodes[self.id].value + rhs,
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }
}

impl<const N: usize> Sub for Variable<N> {
    type Output = Self;
    fn sub(self, rhs: Self) -> Self::Output {
        let mut lhs = self;
        let graph = unsafe { lhs.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::SubX(self.id, rhs.id),
            value: graph.nodes[self.id].value - graph.nodes[rhs.id].value,
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }
}
impl<const N: usize> Sub<f64> for Variable<N> {
    type Output = Self;
    fn sub(self, rhs: f64) -> Self::Output {
        let mut lhs = self;
        let graph = unsafe { lhs.graph.as_mut() };

        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Sub(self.id, rhs),
            value: graph.nodes[self.id].value - rhs,
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }
}

impl<const N: usize> Mul for Variable<N> {
    type Output = Self;
    fn mul(self, rhs: Self) -> Self::Output {
        let mut lhs = self;
        let graph = unsafe { lhs.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::MulX(self.id, rhs.id),
            value: graph.nodes[self.id].value * graph.nodes[rhs.id].value,
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }
}

impl<const N: usize> Mul<f64> for Variable<N> {
    type Output = Self;
    fn mul(self, rhs: f64) -> Self::Output {
        let mut lhs = self;
        let graph = unsafe { lhs.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Mul(self.id, rhs),
            value: graph.nodes[self.id].value * rhs,
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }
}

impl<const N: usize> Div for Variable<N> {
    type Output = Self;
    fn div(self, rhs: Self) -> Self::Output {
        let mut lhs = self;
        let graph = unsafe { lhs.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::DivX(self.id, rhs.id),
            value: graph.nodes[self.id].value / graph.nodes[rhs.id].value,
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }
}

impl<const N: usize> Div<f64> for Variable<N> {
    type Output = Self;
    fn div(self, rhs: f64) -> Self::Output {
        let mut lhs = self;
        let graph = unsafe { lhs.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Div(self.id, rhs),
            value: graph.nodes[self.id].value / rhs,
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }
}

impl<const N: usize> Neg for Variable<N> {
    type Output = Self;
    fn neg(self) -> Self::Output {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Neg(self.id),
            value: -graph.nodes[self.id].value,
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }
}

impl<const N: usize> Variable<N> {
    pub fn recip(self) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Recip(self.id),
            value: graph.nodes[self.id].value.recip(),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn sin(self) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Sin(self.id),
            value: graph.nodes[self.id].value.sin(),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn asin(self) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Asin(self.id),
            value: graph.nodes[self.id].value.asin(),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn sinh(self) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Sinh(self.id),
            value: graph.nodes[self.id].value.sinh(),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn asinh(self) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Asinh(self.id),
            value: graph.nodes[self.id].value.asinh(),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn cos(self) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Cos(self.id),
            value: graph.nodes[self.id].value.cos(),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn acos(self) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Acos(self.id),
            value: graph.nodes[self.id].value.acos(),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn cosh(self) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Cosh(self.id),
            value: graph.nodes[self.id].value.cosh(),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn acosh(self) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Acosh(self.id),
            value: graph.nodes[self.id].value.acosh(),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn tan(self) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Tan(self.id),
            value: graph.nodes[self.id].value.tan(),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn atan(self) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Atan(self.id),
            value: graph.nodes[self.id].value.atan(),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn tanh(self) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Tanh(self.id),
            value: graph.nodes[self.id].value.tanh(),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn atanh(self) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Atanh(self.id),
            value: graph.nodes[self.id].value.atanh(),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn exp(self) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Exp(self.id),
            value: graph.nodes[self.id].value.exp(),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn exp2(self) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Exp2(self.id),
            value: graph.nodes[self.id].value.exp2(),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn exp_m1(self) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::ExpM1(self.id),
            value: graph.nodes[self.id].value.exp_m1(),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn ln(self) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Ln(self.id),
            value: graph.nodes[self.id].value.ln(),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn log10(self) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Log10(self.id),
            value: graph.nodes[self.id].value.log10(),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn log2(self) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Log2(self.id),
            value: graph.nodes[self.id].value.log2(),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn log(self, base: f64) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Log(self.id, base),
            value: graph.nodes[self.id].value.log(base),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn logx(self, base: Self) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::LogX(self.id, base.id),
            value: graph.nodes[self.id].value.log(graph.nodes[base.id].value),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn powf(self, n: f64) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Powf(self.id, n),
            value: graph.nodes[self.id].value.powf(n),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn powi(self, n: i32) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Powi(self.id, n),
            value: graph.nodes[self.id].value.powi(n),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn powx(self, n: Self) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::PowX(self.id, n.id),
            value: graph.nodes[self.id].value.powf(graph.nodes[self.id].value),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn sqrt(self) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Sqrt(self.id),
            value: graph.nodes[self.id].value.sqrt(),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn cbrt(self) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Cbrt(self.id),
            value: graph.nodes[self.id].value.cbrt(),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }

    pub fn hypot(self, other: Self) -> Self {
        let mut var = self;
        let graph = unsafe { var.graph.as_mut() };
        let id = graph.nodes.len();
        let node = Node {
            node_type: NodeType::Hypot(self.id, other.id),
            value: graph.nodes[self.id]
                .value
                .hypot(graph.nodes[other.id].value),
            grad: 0.0,
        };
        graph.nodes.push(node);
        Variable {
            graph: NonNull::from(graph),
            id,
        }
    }
}
